import os
import asyncio
from dotenv import load_dotenv
import streamlit as st
from llm import initialize_llm, get_chain
from prompt import decision_making_agent_prompt, prompt
from ui import display_chat, display_user_message
from ui import sidebar
from utils import highlight_inputs_and_get_metadata_async
from typing import Tuple, Any
# from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_openai import ChatOpenAI  # 导入 ChatOpenAI

load_dotenv(".env")

def init_settings() -> Tuple[asyncio.AbstractEventLoop, ChatOpenAI, Any]:
    loop = asyncio.ProactorEventLoop()
    asyncio.set_event_loop(loop)
    llm = initialize_llm()
    chain = get_chain(prompt, llm)
    return loop, llm, chain

def init_sessions():
    # Initialize session state
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []
        
    if "user_input" not in st.session_state:
        st.session_state.user_input = ""
        
    if "api_key_configured" not in st.session_state:
        st.session_state.api_key_configured = False
        
    if "chain" not in st.session_state:
        st.session_state.chain = chain
        
    if "llm" not in st.session_state:
        st.session_state.llm = llm

# Define the chain as a global variable
chain = None

# Get all the initial settings 
loop, llm, chain = init_settings()


def main():
    # Set page title
    st.set_page_config(page_title="🤖 填表智能体", layout="centered")
    st.title("🤖 填表智能体")

    # Initialize sessions
    init_sessions()

    # Side bar for API key input
    sidebar()

    # Check if API key is configured
    if not st.session_state.api_key_configured and not os.environ.get("OPENAI_API_KEY"):
        st.warning("⚠️ 请在侧边栏配置您的Google API密钥以使用此应用。")

    # Chat input
    if st.session_state.api_key_configured or os.environ.get("OPENAI_API_KEY"):
        user_input = st.chat_input("询问有关表单填写的问题...")
    else:
        user_input = st.chat_input("请先配置API密钥...", disabled=True)

    # Display chat interface
    display_chat()

    if user_input:
        user_msg = {"role": "user", "content": user_input}
        st.session_state.chat_history.append(user_msg)
        
        display_user_message(user_msg)
        
        # Spinner for processing
        with st.spinner("思考中...", show_time=True):
            current_chain = st.session_state.get("chain", chain)
            response = current_chain.invoke({
                "decision_making_agent_prompt": decision_making_agent_prompt,
                "user_prompt": user_input
            })

        if response.get("isItFormFillingRequest"):
            st.session_state.chat_history.append({
                "role": "assistant", 
                "content": response.get("chat", "处理您的请求时遇到问题。")
            })
            url_to_fill = response.get("urlWhereDataToFill")
            data_to_fill = response.get("dataToFill")
            print("Done")
            sreenshot_bytes = loop.run_until_complete(highlight_inputs_and_get_metadata_async(url_to_fill, llm, data_to_fill))
            print("Screenshot captured")
            
        else:    
            st.session_state.chat_history.append({
                "role": "assistant", 
                "content": response.get("chat", "处理您的请求时遇到问题。")
            })
        
        st.rerun()
        
if __name__ == '__main__':
    main()

